#include "precompiled.h"

#include "CBDGenerator.h"

#include "CUtils.h"

#include <windows.h>
#include <iostream>
#include <algorithm>

bool CDBGenerator::generate(const char* dir, const char* base_xml)
{
	TiXmlDocument doc(base_xml);
	if ( !doc.LoadFile() ) doc.InsertEndChild( TiXmlElement("base") );
	TiXmlElement * root = doc.RootElement();

	std::cout << addSignaturesFromDirectory(dir, root) << " signatures have been added to " << base_xml << std::endl;

	return doc.SaveFile();
}

int CDBGenerator::addSignaturesFromDirectory(const char* dir, TiXmlElement* root)
{
	int result = 0;

	char buffer[MAX_PATH];
	GetCurrentDirectory(MAX_PATH, buffer);

	std::string path = std::string(buffer) + "\\" + dir + "\\";
	std::string mask = path + "*.sig";

	WIN32_FIND_DATA FindFileData;

	HANDLE hFile = FindFirstFile(mask.c_str(), &FindFileData);

	if (hFile == INVALID_HANDLE_VALUE) return result;

	BOOL done = FALSE;
	for (; !done; done = !FindNextFile(hFile,&FindFileData))
	{
		if ( !strcmp(FindFileData.cFileName, ".") ) continue;
		if ( !strcmp(FindFileData.cFileName, "..") ) continue;
		if ( !strcmp(FindFileData.cFileName, ".svn") ) continue;

		result += addSignaturesFromFile( (path + FindFileData.cFileName).c_str(), root );
	}

	FindClose(hFile);
	return result;
}

int CDBGenerator::addSignaturesFromFile(const char* filename, TiXmlElement* root)
{
	int result = 0;

	std::vector<std::string> names;
	getNames(root->FirstChildElement("signature"), names);

	FILE* file = fopen(filename, "r");
	if (!file) return result;

	char buffer[4096] = {0};
	while ( fgets(buffer, sizeof(buffer), file) )
	{
		std::string line(buffer);
		line.erase( line.size() - 1 );
		if (line.empty()) continue;

		size_t pos = line.find("=");
		if (pos == std::string::npos) continue;

		std::string name = line.substr(0, pos);
		std::string sig_hex = line.substr(pos + 1);

		if ( std::find( names.begin(), names.end(), name ) != names.end() ) continue;

		TiXmlElement sig("signature");

		sig.SetAttribute("name", name);

		std::string sig_b64;
		CUtils::encodeBase64( CUtils::str2bin(sig_hex.c_str()), sig_b64 );

		sig.SetAttribute("size", sig_b64.size());
		sig.InsertEndChild( TiXmlText(sig_b64) );

		root->InsertEndChild( sig );

		++result;
	}

	fclose(file);
	return result;
}

void CDBGenerator::getNames(const TiXmlElement * first, std::vector<std::string>& names)
{
	if (!first) return;

	names.clear();
	do
	{
		names.push_back( first->Attribute("name") );
	}
	while ( (first = first->NextSiblingElement("signature")) );
}
